#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Ollama + 千问 + MCP 本地 Agent 示例
"""
import json

from qwen_agent.agents import Assistant
from qwen_agent.gui import WebUI
from qwen_agent.tools.base import register_tool, BaseTool
from qwen_agent.utils.output_beautify import typewriter_print

# 1. 把本地 Ollama 包装成千问 Agent 的“大脑”
llm_cfg = {
    'model': 'qwen2.5:1.5b',  # 与 ollama list 里的名字保持一致
    'model_server': 'http://localhost:11434/v1',  # Ollama 默认端口
    'api_key': 'EMPTY'  # Ollama 不需要密钥
}


# 2.1 声明自定义工具
@register_tool('get_weather')
class GetWeather(BaseTool):
    # description的质量好坏直接关系到LLM能否调用到工具，不要乱写
    description = '获取城市当前天气'
    parameters = [{
        'name': 'city',
        'type': 'string',
        'description': '城市名',
        'required': True
    }]

    def call(self, params: str, **kwargs) -> str:
        city = json.loads(params)['city']
        # 这里 mock 结果，真实场景可调用外部 API
        return f"{city} 今天晴，25℃"


# 2.2 声明 MCP 工具（time 获取本地时间，fetch 抓取网页）
# 需用pip提前安装mcp-server-time和mcp-server-fetch
# agent会自动拉起mcp_server_time和mcp_server_fetch两个mcp server，它们均接受json-rpc输入
# mcp server定义中有工具的description，方便agent匹配用户输入
tools = [
    {'mcpServers': {
        'time': {
            'command': 'python',
            'args': ['-m', 'mcp_server_time', '--local-timezone=Asia/Shanghai']
        },
        'fetch': {
            'command': 'python',
            'args': ['-m', 'mcp_server_fetch']
        }
    }},
    'get_weather',
]

# 3. 初始化 Agent
bot = Assistant(llm=llm_cfg, function_list=tools)


# 4. 运行
def run4console():
    messages = []
    while True:
        query = input("\n【用户】>>> ")
        if query.lower() in {"q", "quit"}:
            break

        # 输入查询：
        # 上海天气如何？
        # 抓取 https://qwenlm.github.io/blog/内容 并总结 2 句话
        # 将触发tool calling
        messages.append({'role': 'user', 'content': query})
        resp_plain = ''
        for rsp in bot.run(messages):
            # streaming output
            resp_plain = typewriter_print(rsp, resp_plain)


def run4web():
    WebUI(bot).run()


if __name__ == '__main__':
    run4web()
